import os
import sys
import commands
from swift.common import ring

swift_dir = '/etc/swift'
# TODO: define thread num and batch recover num later


# return ring builder file and ring based on type
def get_target_ring(controller_type):
    ret = None
    if controller_type == "account":
        ret = (os.path.join(swift_dir, "account.builder"), ring.Ring(swift_dir, ring_name="account"))
    elif controller_type == "container":
        ret = (os.path.join(swift_dir, "container.builder"), ring.Ring(swift_dir, ring_name="container"))
    elif controller_type == "obj":
        ret = (os.path.join(swift_dir, "object.builder"), ring.Ring(swift_dir, ring_name="object"))
    else:
        print "invalid controller type"
        sys.exit(1)
    return ret

# get all the partitions that a specific device owns
def get_dev_parts(dev_id, controller_type):
    dev_id = "d"+str(dev_id)
    ring_builder_file = get_target_ring(controller_type)[0]
    cmd = 'swift-ring-builder %s list_parts %s' % (ring_builder_file, dev_id)
    status, output = commands.getstatusoutput(cmd)
    if status:
        print "retrieve partitions for dev %s hit error %s" % (dev_id, output)
        sys.exit(1)

    output = output.split('\n')
    parts = []
    for entry in output[1:]:
        part = entry.strip().split(' ', 1)[0]
        parts.append(part)
    return parts

# get other nodes which occupy the same partition
def get_other_nodes_own_part(dev_id, part, controller_type):
    target_ring = get_target_ring(controller_type)[1]
    nodes = target_ring.get_part_nodes(int(part))
    return [n for n in nodes if n['id'] != dev_id]

# recover all the partitions specified by parts for dev
def recover_parts(dev_id, parts, controller_type):
    if parts == "ALL":
        parts = get_dev_parts(dev_id, controller_type)

    # {ip1: {sdb1 : [parts1], sdc1 : [parts2]}, ip2: {sdb1 : [parts3], sdc1 : [parts4]}}
    recover_info = {}
    for part in parts:
        for node in get_other_nodes_own_part(dev_id, part, controller_type):
            ip = node['ip']
            device = node['device'] # device name
            device_parts = recover_info[ip]
            if device_parts:
                device_parts[device].append(part)
            else:
                recover_info[ip] = {device : [part]}

    # TODO: based on node, ssh each node to perform batch recovery task
    # if parts folder doesn't exist under device, simply skip as it indicates no data under the parts


if __name__ == "__main__":
    # TODO: add options
    # part 0,3 of d0 are missing, trigger recovery which let its peers (nodes also own 0,3 parts) do replication.
    recover_parts(0, [0,3], "container")